# -*- coding: utf-8 -*-#
"""
@File : 03-余弦相似度和欧几里得距离.py
@Description : 
@Author : Le.Qing
@Create Time : 2025-05-08 10:28
"""
import math

from langchain_ollama import OllamaEmbeddings
from numpy import dot
from numpy.linalg import norm


def embeddingLLM():
    return OllamaEmbeddings(base_url="http://localhost:11434", model="chevalblanc/acge_text_embedding")


model = embeddingLLM()
input_texts = [
    "中国的首都是哪里",
    "你喜欢去哪里旅游",
    "北京",
    "今天中午吃什么"
]
embeddings = model.embed_documents(input_texts)

a = embeddings[0]
b = embeddings[2]


def cos_a_b(vec1, vec2):  # 计算余弦相似度
    return dot(vec1, vec2) / (norm(vec1) * norm(vec2))


def euclidean_distance(vec1, vec2):  # 计算欧几里得距离
    if len(vec1) != len(vec2):
        raise ValueError("向量维度必须相同")
    squared_distance = sum((a - b) ** 2 for a, b in zip(vec1, vec2))
    return math.sqrt(squared_distance)


print('余弦相似度对比：(越大越相似)')
for i in range(4):
    print(input_texts[0], input_texts[i], cos_a_b(embeddings[0], embeddings[i]))
print('欧几里得距离对比：(越小越近)')
for i in range(4):
    print(input_texts[0], input_texts[i], euclidean_distance(embeddings[0], embeddings[i]))
